from PyQt5.QtWidgets import *
from PIL import Image
from ui.WinUI import Ui_MainWindow
from itertools import product
import fitz
import sys,os
import shutil

class MyWindow(QWidget,Ui_MainWindow):
    def __init__(self):
        super(MyWindow, self).__init__()
        self.setupUi(self)
        self.openbtn.clicked.connect(self.openFile)
        self.savebtn.clicked.connect(self.savefile)


    def openFile(self):
        if self.color_rgb.text()!="":
            filename,filetype = QFileDialog.getOpenFileName(self,"open file",'.','Image Files(*.jpg *.png);;PDF Files(*.pdf)')
            self.openfilename.setText(filename)
            if filetype == "Image Files(*.jpg *.png)":
                self.img = Image.open(filename)
                width, height = self.img.size
                for row in range(width):
                    for col in range(height):
                        pos =  (row,col)
                        rgb = self.img.getpixel(pos)[:3]
                        if (sum(rgb) >= int(self.color_rgb.text())):
                            self.img.putpixel(pos, (255, 255, 255))
            else:
                self.remove_pdf()
        else:
            QMessageBox.information(self,
            "温馨提示",
            "请先输入颜色阈值（RGB三色值之和）")

    def savefile(self):
        outfilename,outfiletype = QFileDialog.getSaveFileName(self,'save file','.','Image Files(*.jpg);;Image Files(*.png);;PDF Files(*.pdf)')
        self.savefilename.setText(outfilename)
        if outfiletype == "Image Files(*.jpg)":
            self.img = self.img.convert('RGB')
        if outfiletype == "Image Files(*.jpg)" or outfiletype == "Image Files(*.png)" :
            self.img.save(outfilename)
        else:
            self.pic2pdf()


    def remove_pdf(self):
        QMessageBox.information(self, '提示', '请等待操作。。')
        page_num = 0
        pdf_file = self.openfilename.text()
        if pdf_file != "":
            path = os.path.join(os.getcwd(),"pdf_images")
            folder = os.path.exists(path)
            if not folder:  # 判断是否存在文件夹如果不存在则创建为文件夹
                os.makedirs(path)  # makedirs 创建文件时如果路径不存在会创建这个路径
            else:
                shutil.rmtree(path)
                os.makedirs(path)
            pdf = fitz.open(pdf_file);
            for page in pdf:
                pixmap = page.get_pixmap()
                for pos in product(range(pixmap.width), range(pixmap.height)):
                    rgb = pixmap.pixel(pos[0], pos[1])
                    if(sum(rgb) >= int(self.color_rgb.text())):
                        pixmap.set_pixel(pos[0], pos[1], (255, 255, 255))
                pixmap.pil_save(f"pdf_images/{page_num}.png")
                print(f"第{page_num}水印去除完成")
                page_num = page_num + 1
            QMessageBox.information(self,'提示','操作成功，请保存新文件。')
        else:
            QMessageBox.information(self, '提示', '请选择文件')
    def pic2pdf(self):
        pic_dir = "pdf_images"
        pdf = fitz.open()
        img_files = sorted(os.listdir(pic_dir), key=lambda x: int(str(x).split('.')[0]))
        for img in img_files:
            print(img)
            imgdoc = fitz.open(pic_dir + '/' + img)
            pdfbytes = imgdoc.convert_to_pdf()
            imgpdf = fitz.open("pdf", pdfbytes)
            pdf.insert_pdf(imgpdf)
        pdf.save(self.savefilename.text())
        pdf.close()

if __name__ == '__main__':
    app = QApplication(sys.argv)
    win = MyWindow()
    win.show()
    sys.exit(app.exec_())